package utils

import (
	"bytes"
	"crypto/aes"
)

func AESEncryptECBToBase64(plaintext string, key string) (string, error) {
	b, err := AESEncryptECB([]byte(plaintext), []byte(key))
	if err != nil {
		return "", err
	}
	return Base64Encode(b), nil
}

func AESDecryptECBFromBase64(base64text string, key string) ([]byte, error) {
	ciphertext, err := Base64Decode(base64text)
	if err != nil {
		return nil, err
	}
	return AESDecryptECB(ciphertext, []byte(key))
}

// AESEncryptECB 使用 AES ECB 模式加密数据
func AESEncryptECB(plaintext, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 分组填充方式为 PKCS7Padding
	plaintext = PKCS7Padding(plaintext, block.BlockSize())

	ciphertext := make([]byte, len(plaintext))

	// 每个分组使用相同的密钥进行加密
	for i := 0; i < len(plaintext); i += block.BlockSize() {
		block.Encrypt(ciphertext[i:i+block.BlockSize()], plaintext[i:i+block.BlockSize()])
	}

	return ciphertext, nil
}

// AESDecryptECB 使用 AES ECB 模式解密数据
func AESDecryptECB(ciphertext, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	plaintext := make([]byte, len(ciphertext))

	// 每个分组使用相同的密钥进行解密
	for i := 0; i < len(ciphertext); i += block.BlockSize() {
		block.Decrypt(plaintext[i:i+block.BlockSize()], ciphertext[i:i+block.BlockSize()])
	}

	// 去除填充数据
	plaintext = PKCS7UnPadding(plaintext)

	return plaintext, nil
}

// PKCS7Padding 使用 PKCS7 填充方式对数据进行填充
func PKCS7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(data, padText...)
}

// PKCS7UnPadding 使用 PKCS7 填充方式对数据进行去填充
func PKCS7UnPadding(data []byte) []byte {
	length := len(data)
	unPadding := int(data[length-1])
	return data[:(length - unPadding)]
}
